import os
import re
import tkinter as tk
from tkinter import filedialog, messagebox, ttk
from openai import OpenAI
import threading


class LatexTranslatorApp:
    def __init__(self, root):
        self.root = root
        self.root.title("LaTeX Translator")
        self.root.geometry("1200x800")

        self.tex_file_path = None
        self.tex_content = None
        self.sections = []
        self.current_section = None
        self.translated_content = None

        self.create_ui()

    def create_ui(self):
        # Main frame
        main_frame = ttk.Frame(self.root, padding=10)
        main_frame.pack(fill=tk.BOTH, expand=True)

        # File selection
        file_frame = ttk.Frame(main_frame)
        file_frame.pack(fill=tk.X, pady=5)

        ttk.Label(file_frame, text="LaTeX File:").pack(side=tk.LEFT, padx=5)
        self.file_entry = ttk.Entry(file_frame, width=50)
        self.file_entry.pack(side=tk.LEFT, padx=5, fill=tk.X, expand=True)
        ttk.Button(file_frame, text="Browse", command=self.select_file).pack(side=tk.LEFT, padx=5)

        # Action buttons at the top
        button_frame = ttk.Frame(main_frame)
        button_frame.pack(fill=tk.X, pady=5, padx=5)

        # Using grid for better control over button layout
        translate_btn = ttk.Button(button_frame, text="Translate", command=self.start_translation, width=20)
        translate_btn.grid(row=0, column=0, padx=10, pady=5, sticky="ew")

        replace_btn = ttk.Button(button_frame, text="Replace in Original", command=self.replace_section, width=20)
        replace_btn.grid(row=0, column=1, padx=10, pady=5, sticky="ew")

        save_btn = ttk.Button(button_frame, text="Save to New File", command=self.save_file, width=20)
        save_btn.grid(row=0, column=2, padx=10, pady=5, sticky="ew")

        # Configure the grid to distribute extra space
        button_frame.columnconfigure(0, weight=1)
        button_frame.columnconfigure(1, weight=1)
        button_frame.columnconfigure(2, weight=1)

        # Split view
        paned_window = ttk.PanedWindow(main_frame, orient=tk.HORIZONTAL)
        paned_window.pack(fill=tk.BOTH, expand=True, pady=10)

        # Left panel - Section list
        left_frame = ttk.Frame(paned_window)
        paned_window.add(left_frame, weight=1)

        ttk.Label(left_frame, text="Sections:").pack(anchor=tk.W)

        # Add a scrollbar to the treeview
        tree_scroll = ttk.Scrollbar(left_frame)
        tree_scroll.pack(side=tk.RIGHT, fill=tk.Y)

        self.section_tree = ttk.Treeview(left_frame, yscrollcommand=tree_scroll.set)
        tree_scroll.config(command=self.section_tree.yview)

        self.section_tree["columns"] = ("type", "title")
        self.section_tree.column("#0", width=0, stretch=tk.NO)
        self.section_tree.column("type", width=100, anchor=tk.W)
        self.section_tree.column("title", width=200, anchor=tk.W)
        self.section_tree.heading("type", text="Type")
        self.section_tree.heading("title", text="Title")
        self.section_tree.pack(fill=tk.BOTH, expand=True)
        self.section_tree.bind("<<TreeviewSelect>>", self.on_section_select)

        # Right panel - Content and translation
        right_frame = ttk.Frame(paned_window)
        paned_window.add(right_frame, weight=2)

        right_paned = ttk.PanedWindow(right_frame, orient=tk.VERTICAL)
        right_paned.pack(fill=tk.BOTH, expand=True)

        # Original content
        original_frame = ttk.LabelFrame(right_paned, text="Original LaTeX")
        right_paned.add(original_frame, weight=1)

        # Add scrollbar to original text
        orig_scroll_y = ttk.Scrollbar(original_frame)
        orig_scroll_y.pack(side=tk.RIGHT, fill=tk.Y)
        orig_scroll_x = ttk.Scrollbar(original_frame, orient=tk.HORIZONTAL)
        orig_scroll_x.pack(side=tk.BOTTOM, fill=tk.X)

        self.original_text = tk.Text(original_frame, wrap=tk.NONE,
                                     yscrollcommand=orig_scroll_y.set,
                                     xscrollcommand=orig_scroll_x.set)
        self.original_text.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        orig_scroll_y.config(command=self.original_text.yview)
        orig_scroll_x.config(command=self.original_text.xview)

        # Translation
        translation_frame = ttk.LabelFrame(right_paned, text="Translation")
        right_paned.add(translation_frame, weight=1)

        # Add scrollbar to translation text
        trans_scroll_y = ttk.Scrollbar(translation_frame)
        trans_scroll_y.pack(side=tk.RIGHT, fill=tk.Y)
        trans_scroll_x = ttk.Scrollbar(translation_frame, orient=tk.HORIZONTAL)
        trans_scroll_x.pack(side=tk.BOTTOM, fill=tk.X)

        self.translation_text = tk.Text(translation_frame, wrap=tk.NONE,
                                        yscrollcommand=trans_scroll_y.set,
                                        xscrollcommand=trans_scroll_x.set)
        self.translation_text.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        trans_scroll_y.config(command=self.translation_text.yview)
        trans_scroll_x.config(command=self.translation_text.xview)

        # Status bar and progress bar are now at the bottom

        # Status bar
        self.status_var = tk.StringVar()
        self.status_var.set("Ready")
        status_bar = ttk.Label(main_frame, textvariable=self.status_var, relief=tk.SUNKEN, anchor=tk.W)
        status_bar.pack(fill=tk.X, side=tk.BOTTOM, pady=5)

        # Progress bar
        self.progress_bar = ttk.Progressbar(main_frame, mode='indeterminate')
        self.progress_bar.pack(fill=tk.X, side=tk.BOTTOM, pady=5)

    def select_file(self):
        file_path = filedialog.askopenfilename(filetypes=[("LaTeX files", "*.tex"), ("All files", "*.*")])
        if file_path:
            self.tex_file_path = file_path
            self.file_entry.delete(0, tk.END)
            self.file_entry.insert(0, file_path)
            self.load_tex_file()

    def load_tex_file(self):
        try:
            self.status_var.set("Loading LaTeX file...")
            self.progress_bar.start()

            with open(self.tex_file_path, 'r', encoding='utf-8') as file:
                self.tex_content = file.read()

            # Extract sections and subsections
            self.extract_sections()
            self.populate_section_tree()

            self.status_var.set("LaTeX file loaded successfully.")
            self.progress_bar.stop()
            messagebox.showinfo("Success", "LaTeX file loaded successfully.")
        except Exception as e:
            self.status_var.set(f"Error: {str(e)}")
            self.progress_bar.stop()
            messagebox.showerror("Error", f"Failed to load LaTeX file: {str(e)}")

    def extract_sections(self):
        self.sections = []

        # This regex matches sections with their entire content
        section_pattern = r'(\\section\{([^}]*)\})(.*?)(?=\\section\{|$)'

        # Find all sections
        section_matches = re.finditer(section_pattern, self.tex_content, re.DOTALL)
        for i, match in enumerate(section_matches):
            full_section = match.group(0)
            section_cmd = match.group(1)
            section_title = match.group(2)
            section_content = match.group(3)
            section_start = match.start()
            section_end = match.end()

            section_id = f"section_{i}"
            self.sections.append({
                'id': section_id,
                'type': 'section',
                'title': section_title,
                'full_content': full_section,
                'content': section_content,
                'start': section_start,
                'end': section_end,
                'subsections': []
            })

            # Find all subsections within this section
            subsection_pattern = r'(\\subsection\{([^}]*)\})(.*?)(?=\\subsection\{|$)'
            subsection_matches = re.finditer(subsection_pattern, section_content, re.DOTALL)

            for j, sub_match in enumerate(subsection_matches):
                full_subsection = sub_match.group(0)
                subsection_cmd = sub_match.group(1)
                subsection_title = sub_match.group(2)
                subsection_content = sub_match.group(3)

                # Calculate the absolute position in the original text
                rel_start = sub_match.start()
                rel_end = sub_match.end()
                abs_start = section_start + len(section_cmd) + rel_start
                abs_end = section_start + len(section_cmd) + rel_end

                subsection_id = f"{section_id}_subsection_{j}"
                self.sections[-1]['subsections'].append({
                    'id': subsection_id,
                    'type': 'subsection',
                    'title': subsection_title,
                    'full_content': full_subsection,
                    'content': subsection_content,
                    'start': abs_start,
                    'end': abs_end
                })

    def populate_section_tree(self):
        self.section_tree.delete(*self.section_tree.get_children())
        for section in self.sections:
            section_item = self.section_tree.insert("", "end", text="", values=("section", section['title']),
                                                    iid=section['id'], open=True)

            for subsection in section['subsections']:
                self.section_tree.insert(section_item, "end", text="", values=("subsection", subsection['title']),
                                         iid=subsection['id'])

    def on_section_select(self, event):
        selected_items = self.section_tree.selection()
        if not selected_items:
            return

        selected_item = selected_items[0]

        # Determine if this is a section or subsection
        if "_subsection_" in selected_item:
            # Find the subsection
            section_id, subsection_idx = selected_item.split("_subsection_")
            section_idx = int(section_id.split("_")[1])
            subsection_idx = int(subsection_idx)

            self.current_section = self.sections[section_idx]['subsections'][subsection_idx]
        else:
            # It's a main section
            section_idx = int(selected_item.split("_")[1])
            self.current_section = self.sections[section_idx]

        self.original_text.delete(1.0, tk.END)
        self.original_text.insert(1.0, self.current_section['full_content'])
        self.translation_text.delete(1.0, tk.END)
        self.status_var.set(f"Loaded {self.current_section['type']}: {self.current_section['title']}")

    def start_translation(self):
        if not self.current_section:
            messagebox.showwarning("Warning", "Please select a section to translate.")
            return

        # Start translation in a separate thread to keep UI responsive
        self.progress_bar.start()
        self.status_var.set("Translating...")

        threading.Thread(target=self.translate_section, daemon=True).start()

    def translate_section(self):
        try:
            content_to_translate = self.original_text.get(1.0, tk.END)

            # Call OpenAI API
            client = OpenAI(
                api_key="0f56e8eb-8815-482d-8099-c5bef02f129a",
                base_url="https://ark.cn-beijing.volces.com/api/v3",
            )

            # The translation prompt for the AI
            translation_prompt = """你是一个专业的翻译助手，负责将英文LaTeX文档的section内容翻译成中文。请根据以下要求完成任务：
1. 翻译时要保留原始LaTeX文档的结构、命令、标签和其他格式元素。
2. 请仅翻译文本内容，不要改变任何LaTeX命令或环境格式（例如：\\section{Title}、\\textbf{bold text}等）。
3. LaTeX命令中的参数需要翻译，例如\\section{Introduction}应翻译为\\section{引言}。
4. 确保翻译后的中文句子与原文的语法和意思一致，并自然流畅。
5. 不要更改LaTeX代码中用于排版、引用等的任何命令。
6. 请将翻译后的内容直接返回，并保持原文的LaTeX格式不变。
7. 不要在回复中包含额外的注释或解释，只需要返回翻译后的内容。
8. 你收到的内容是一个完整的section，请确保整体翻译，不要遗漏任何部分。"""

            completion = client.chat.completions.create(
                model="deepseek-r1-250120",
                messages=[
                    {"role": "system", "content": translation_prompt},
                    {"role": "user", "content": content_to_translate}
                ],
            )

            self.translated_content = completion.choices[0].message.content

            # Update UI in the main thread
            self.root.after(0, self.update_translation_ui, self.translated_content)

        except Exception as e:
            self.root.after(0, self.show_translation_error, str(e))

    def update_translation_ui(self, translated_text):
        self.translation_text.delete(1.0, tk.END)
        self.translation_text.insert(1.0, translated_text)
        self.progress_bar.stop()
        self.status_var.set("Translation complete")

    def show_translation_error(self, error_msg):
        self.progress_bar.stop()
        self.status_var.set(f"Translation error: {error_msg}")
        messagebox.showerror("Translation Error", f"Failed to translate: {error_msg}")

    def replace_section(self):
        if not self.current_section or not self.translated_content:
            messagebox.showwarning("Warning", "Please select a section and translate it first.")
            return

        try:
            # Replace the section content in the original text
            start = self.current_section['start']
            end = self.current_section['end']

            # Ensure we're using the translated_content from the text widget
            # in case the user made manual edits
            translated_text = self.translation_text.get(1.0, tk.END).strip()

            new_tex_content = self.tex_content[:start] + translated_text + self.tex_content[end:]
            self.tex_content = new_tex_content

            # Re-extract sections to update the data
            self.extract_sections()
            self.populate_section_tree()

            self.status_var.set("Section replaced successfully")
            messagebox.showinfo("Success", "Section replaced successfully.")
        except Exception as e:
            self.status_var.set(f"Error replacing section: {str(e)}")
            messagebox.showerror("Error", f"Failed to replace section: {str(e)}")

    def save_file(self):
        if not self.tex_content:
            messagebox.showwarning("Warning", "No LaTeX content to save.")
            return

        file_path = filedialog.asksaveasfilename(
            defaultextension=".tex",
            filetypes=[("LaTeX files", "*.tex"), ("All files", "*.*")]
        )

        if file_path:
            try:
                with open(file_path, 'w', encoding='utf-8') as file:
                    file.write(self.tex_content)
                self.status_var.set(f"File saved to {file_path}")
                messagebox.showinfo("Success", f"File saved to {file_path}")
            except Exception as e:
                self.status_var.set(f"Error saving file: {str(e)}")
                messagebox.showerror("Error", f"Failed to save file: {str(e)}")


if __name__ == "__main__":
    root = tk.Tk()
    app = LatexTranslatorApp(root)
    root.mainloop()